from vgg13_shortcuts import vgg13_shortcuts_v2
import os

def main():
    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    model_path_ori = 'models_ori'
    model_path_out = 'models'
    weight_decay = 1e-5
    num_classes = (2,)
    feature_names = ('Mal_2',)
    model = vgg13_shortcuts_v2((64, 64, 64, 1), 32, weight_decay, num_classes, feature_names)

    model_names = os.listdir(model_path_ori)
    for model_name in model_names:
        model_epoches = os.listdir(os.path.join(model_path_ori, model_name))
        model_epoches = [x for x in model_epoches if x.startswith('epoch')]
        for model_epoch in model_epoches:
            model_indir = os.path.join(model_path_ori, model_name, model_epoch)
            model_outdir = os.path.join(model_path_out, model_name, model_epoch)
            print('processing ', model_indir)
            os.makedirs(os.path.join(model_path_out, model_name), exist_ok=True)
            model.load_weights(model_indir, by_name=True)
            model.save_weights(model_outdir, overwrite=False)




if __name__ == '__main__':
    main()
